#include "math.h"
#include "stdlib.h"
#include "stdio.h"
#include "string.h"
#include "inc/core/ucgrid1d.h"
#include "gsl/gsl_errno.h"
#include "gsl/gsl_randist.h"


#define PARTICLE_IN_CELL 1
#define PARTICLE_OUT_CELL 2

/*utilites*/
void linear_interpolation(floatdig *l_val, floatdig* l_pos,
        floatdig *r_val, floatdig* r_pos,
        floatdig *pos, floatdig* interpolated_val) {
    floatdig k = (*l_val - *r_val) / (*l_pos - *r_pos);
    floatdig b = *l_val - k * (*l_pos);
    *interpolated_val = k * (*pos) + b;
}

int is_particle_in_cell(floatdig pos, floatdig lbound, floatdig rbound) {
    if ((pos < rbound) && (pos > lbound)) {
        return PARTICLE_IN_CELL;
    } else {
        return PARTICLE_OUT_CELL;
    }
}

int get_num_of_particles_in_cell(ucgrid1d* init, particles_set* set, int cell_num);

/*API*/
void ucgrid1d_init_grid(init_ucgrid1d* init, ucgrid1d* grid) {
#ifdef VERBOSE
    printf("init ucgrid1d\n");
#endif
    if ((init != NULL) && (grid != NULL)) {
       
        floatdig min; 
        floatdig max; 
        int num_of_cells;
        int i = 0;
		memcpy(&(grid->parameters), init, sizeof (init_ucgrid1d));
		min = grid->parameters.min_z;
        max = grid->parameters.max_z;
		
		num_of_cells = ceil((max - min) / (grid->parameters.cell_size)) + 1;
		
		

        //the module uses normalized grid
        //bounds are set to [0,1]
        //corresponded changes for particles' positions and speeds are required 
        //either in mover1d or task1d module
        grid->grid_cell_size = 1. / num_of_cells;
        grid->num_of_cells = num_of_cells;
        
        grid->rho = (floatdig*) malloc(sizeof (floatdig) * (grid->num_of_cells) );
        grid->E = (floatdig*) malloc(sizeof (floatdig) * (grid->num_of_cells ));
        grid->B = (floatdig*) malloc(sizeof (floatdig) * (grid->num_of_cells ));
        grid->phi = (floatdig*) malloc(sizeof (floatdig) * (grid->num_of_cells ));
        grid->cellnums = (floatdig*) malloc(sizeof (floatdig)*(grid->num_of_cells ));
        poisson1d_init_fft(grid->num_of_cells , grid->parameters.cell_size, &(grid->poisson_solver));
        for (i = 0; i < grid->num_of_cells; i++) {
            grid->rho[i] = 0;
            grid->E[i] = 0;
            grid->phi[i] = 0;
            grid->cellnums[i] = i;
        }
    }
}

void ucgrid1d_clear_grid(ucgrid1d* grid) {
    if (grid != NULL) {
        free(grid->rho);
        free(grid->E);
        free(grid->phi);
        free(grid->B);
        free(grid->cellnums);
    }
}

void ucgrid1d_calc_charge_density_from_particles_distr(particles_set* set, ucgrid1d* grid) {
    if ((set != NULL) && (grid != NULL)) {
        int num_of_gcells = grid->num_of_cells;
        int i = 0;
        floatdig qc = set->charge / grid->parameters.cell_size;

#ifdef DEBUG
        printf("DEBUG: Calc charge density from prticles distribution\n");
#endif
        for (i = 0; i < num_of_gcells; i++) {
            grid->rho[i] = 0;
        }

        for(i=0; i< set->num_of_particles; i++)
        {
            int lc = 0;
            int rc = 0;
            floatdig lc_pos = 0;
            floatdig rc_pos = 0;
            floatdig p_pos = set->particles[i].pos.v[Z];
            floatdig lq;
            floatdig rq;
            ucgrid1d_search_neighbor_cells(grid, p_pos, &lc, &rc);
            ucgrid1d_get_cell_center_pos(grid, lc, &lc_pos);
            ucgrid1d_get_cell_center_pos(grid, rc, &rc_pos);
            if (lc == qc)
            {
                lq = qc;
                rq = 0.0;
            }
            else
            {
                lq = qc*(rc_pos - p_pos)/grid->parameters.cell_size;
                rq = qc*(p_pos - lc_pos)/grid->parameters.cell_size;
            }
            grid->rho[lc] += lq;
            grid->rho[rc] += rq;

            if (lc == 0)
            {
                 grid->rho[grid->num_of_cells - 1] += lq;
            }
            if(rc == (grid->num_of_cells - 1))
            {
                grid->rho[0] +=  rq;
            }
            if ((rc > grid->num_of_cells - 1) || (lc < 0))
            {
                printf("ALARM\n");
            }
        }
    }
    
}

int ucgrid1d_get_cell_boundaries(ucgrid1d* grid, int cell_num, floatdig * left, floatdig * right) {
    floatdig cell_size = grid->parameters.cell_size;
    int num_of_cells = grid->num_of_cells;

    if (cell_num == 0) {
        *right = 0.5*cell_size;
        *left = 0;
        return CELLTYPE_LEFTBOUND;
    } else if (cell_num == num_of_cells) {
        *left = cell_size * (cell_num - 0.5);
        *right = cell_size * cell_num;
        return CELLTYPE_RIGHTBOUND;
    } else if (cell_num > grid->num_of_cells) {
        return ERROR;
    } else {
        *right = cell_size * (cell_num + 0.5);
        *left = cell_size *  (cell_num - 0.5);
        return CELLTYPE_REGULAR;
    }
}

int get_num_of_particles_in_cell(ucgrid1d* grid, particles_set* set, int cell_num) {

    int num_of_particles = set->num_of_particles;
    int i = 0;
    floatdig left;
    floatdig right;
    particle* particles = set->particles;
    int num_of_particles_in_cell = 0;
    int celltype = ucgrid1d_get_cell_boundaries(grid, cell_num, &left, &right);

    //it is supposed the function gets normilized to grid's cell size
    //particles' coordinates


    for (i = 0; i < num_of_particles; i++) {
        if ((particles[i].pos.v[Z] > left) && (particles[i].pos.v[Z] < right) && (particles[i].active == PARTICLE_ACTIVE)) {
            num_of_particles_in_cell++;
        }


    }
    return num_of_particles_in_cell;
}

void ucgrid1d_get_cell_center_pos(ucgrid1d* grid, int cell_num, floatdig* pos) {
    //do not use any checks because we want speeed!
    *pos = grid->parameters.cell_size * cell_num;
}


#ifdef DEBUG

int ucgrid1d_distribute_phi_debug(ucgrid1d* grid) {
    floatdig sL = grid->parameters.cell_size * (grid->num_of_cells - 1);
    floatdig A = 1e-4;
    int i;
    for (i = 0; i < grid->num_of_cells; i++) {
        floatdig x = 0;
        ucgrid1d_get_cell_center_pos(grid, i, &x);
        grid->phi[i] = A * sin((PI / sL) * x) * sin((PI / sL) * x);
    }
	return EXIT_SUCCESS;
}

void ucgrid1d_distribute_rho_debug(particles_set* set, ucgrid1d* grid) {
    int num_of_cells = grid->num_of_cells;
    int i = 0;
    floatdig sL = grid->parameters.cell_size * (grid->num_of_cells - 1);
    for (i = 0; i <= num_of_cells / 2; i++) {
        grid->rho[i] = i / sL;
    }
    for (i = num_of_cells / 2; i < num_of_cells; i++) {
        grid->rho[i] = -i / sL + num_of_cells / (sL);
    }
}

int ucgrid1d_distribute_particles_debug(particles_set* set, ucgrid1d* grid) {
    int num_of_particles = set->num_of_particles;
    int num_of_cells = grid->num_of_cells;
    int i = 0;
    int j;
    floatdig A = 0;
    int proc_particles = 0;
    gsl_rng * r;
    gsl_rng_env_setup();
    // select random number generator
    r = gsl_rng_alloc(gsl_rng_mt19937);

#ifdef DEBUG
    printf("DEBUG: Distribute particles for debuf purposes\n");
#endif

    for (i = 0; i < num_of_cells; i++) {
		int particles_to_add = floor((floatdig) num_of_particles / ((floatdig) num_of_cells ));
        floatdig pos = 0;
        ucgrid1d_get_cell_center_pos(grid, i, &pos);      
        j = 0;
#ifdef DEBUG
        printf("DEBUG: Set to cell %d n=%d particles\n", i, particles_to_add);
#endif
        while (j < particles_to_add) {
            floatdig shift = (floatdig) (grid->parameters.cell_size * gsl_rng_uniform(r) - 0.5 * grid->parameters.cell_size);
            
            set->particles[j + proc_particles].pos.v[Z] = pos + shift;
            set->particles[j + proc_particles].active = PARTICLE_ACTIVE;
            j++;
        }
        proc_particles += j;
    }
    return proc_particles;
}

void ucgrid1d_print_rho_into_file_debug(char* filename, ucgrid1d* grid) {
    FILE* to;
    int i = 0;
    to = fopen(filename, "w");
    for (i = 0; i < grid->num_of_cells; i++) {
        fprintf(to, "%e\n", grid->rho[i]);
    }
    fclose(to);
}

void ucgrid1d_print_phi_into_file_debug(char* filename, ucgrid1d* grid) {
    FILE* to;
    int i = 0;
    to = fopen(filename, "w");
    for (i = 0; i < grid->num_of_cells; i++) {
        fprintf(to, "%e\n", grid->phi[i]);
    }
    fclose(to);
}

void ucgrid1d_print_E_into_file_debug(char* filename, ucgrid1d* grid) {
    FILE* to;
    int i = 0;
    to = fopen(filename, "w");
    for (i = 0; i < grid->num_of_cells; i++) {
        fprintf(to, "%e\n", grid->E[i]);
    }
    fclose(to);
}
#endif

void ucgrid1d_calc_phi_from_charge_density(ucgrid1d* grid) {
    floatdig *phi = grid->phi;
    floatdig *rho = grid->rho;
    poisson1d* p1d = &(grid->poisson_solver);
    poisson1d_load_rho(p1d, rho);
    poisson1d_make_fftrho_from_real_rho(p1d);
    //poisson1d_cut_fftrho_harmonics(p1d, 40);
    poisson1d_make_fftphi_from_fftrho(p1d);
    poisson1d_make_phi_from_fftphi(p1d);
    poisson1d_unload_phi(p1d, phi);
}

void ucgrid1d_calc_fields_from_charge_density(ucgrid1d* grid) {
    int i = 0;
    floatdig dx_double = 2 * (grid->parameters.cell_size);

    for (i = 1; i < grid->num_of_cells - 1; i++) {
        grid->E[i] = ((grid->phi[i - 1])-(grid->phi[i + 1])) / dx_double;
    }
    grid->E[0] = -(grid->phi[1] - grid->phi[grid->num_of_cells - 1]) / dx_double; //periodical bonds
    grid->E[(grid->num_of_cells) - 1] = (grid->phi[(grid->num_of_cells) - 2] - grid->phi[0])/dx_double; //periodical bonds
}

int ucgrid1d_search_neighbor_cells(ucgrid1d* grid, floatdig pos, int *left, int *right) {
    //it is supposed the function gets normilized to grid's cell size
    //particles' coordinates

    floatdig a = pos/(grid->parameters.cell_size);
    *left = floor(a);
    *right = ceil(a);
    return OK;
}

void ucgrid1d_interpolate_fields_in_point(ucgrid1d* grid, floatdig pos, floatdig *iefield, floatdig *ibfield) {
    int lc = 0;
    int rc = 0;
    floatdig lc_pos = 0;
    floatdig rc_pos = 0;
    floatdig lE = 0;
    floatdig rE = 0;
    floatdig lB = 0;
    floatdig rB = 0;
    ucgrid1d_search_neighbor_cells(grid, pos, &lc, &rc);
    ucgrid1d_get_cell_center_pos(grid, lc, &lc_pos);
    ucgrid1d_get_cell_center_pos(grid, rc, &rc_pos);
    lE = grid->E[lc];
    rE = grid->E[rc];
    lB = grid->B[lc];
    rB = grid->B[rc];
    linear_interpolation(&lE, &lc_pos, &rE, &rc_pos, &pos, iefield);
    linear_interpolation(&lB, &lc_pos, &rB, &rc_pos, &pos, ibfield);
}

void ucgrid1d_put_estatic_energy_into_file(char* homedir, char* filename, ucgrid1d* grid)
{
    FILE* to;
    char fullname[MAX_BUF_SIZE];
    int k;
    floatdig ES = 0.0;
    sprintf(fullname, "%s/%s", homedir, filename);

    to = fopen(fullname, "w");
    
    for (k = 0; k < grid->poisson_solver.size; k++)
    {
        ES += REAL(grid->poisson_solver.fftphi, k)*REAL(grid->poisson_solver.fftrho, k) + 
              IMAGE(grid->poisson_solver.fftphi, k)*IMAGE(grid->poisson_solver.fftrho, k);
    }
    fprintf(to, "%e\n", ES);
    fclose(to);
}
